import streamlit as st
import os
import glob
import pandas as pd
from PIL import Image
import time
from datetime import datetime
from typing import List

from models.clip_model import ClipModel
from models.kb_factory import KnowledgeBaseFactory
from utils.helpers import get_products_from_data_dir, get_defect_types_from_train_dir, save_upload_file

def show_knowledge_base():
    """显示知识库管理页面"""
    
    st.title("知识库管理")
    
    # 初始化会话状态
    if "clip_model" not in st.session_state:
        st.session_state.clip_model = None
    
    if "knowledge_base" not in st.session_state:
        st.session_state.knowledge_base = None
    
    # CLIP模型选择
    st.header("1. 选择CLIP模型")
    
    # 获取可用的CLIP模型
    available_models = ["ViT-B-16", "ViT-L-14", "ViT-L-14-336", "ViT-H-14"]
    selected_model = st.selectbox("选择CN_CLIP模型", available_models)
    
    # Milvus配置
    st.header("2. 配置向量存储")
    
    milvus_mode = st.radio(
        "选择Milvus模式",
        ["Milvus Lite (本地存储)", "Milvus 服务器"],
        horizontal=True
    )
    
    if milvus_mode == "Milvus 服务器":
        col1, col2 = st.columns(2)
        with col1:
            milvus_host = st.text_input("Milvus 服务器地址", value="localhost")
        with col2:
            milvus_port = st.text_input("Milvus 服务器端口", value="19530")
    else:
        milvus_host = "lite"
        milvus_port = "19530"
    
    # 加载模型按钮
    if st.button("加载模型", key="load_model"):
        with st.spinner(f"加载 {selected_model} 模型中..."):
            try:
                # 初始化CLIP模型
                st.session_state.clip_model = ClipModel(model_name=selected_model)
                st.success(f"{selected_model} 模型加载成功！")
                
                # 初始化知识库
                st.session_state.knowledge_base = KnowledgeBaseFactory.create_knowledge_base(
                    base_path="./knowledge_base",
                    clip_model=st.session_state.clip_model,
                    use_lite=(milvus_mode == "Milvus Lite (本地存储)"),
                    milvus_host=milvus_host,
                    milvus_port=milvus_port
                )
                
                st.success(f"知识库初始化成功: {milvus_mode}")
            except Exception as e:
                st.error(f"模型加载失败: {str(e)}")
    
    # 当前模型信息
    if st.session_state.clip_model is not None:
        st.success(f"当前模型: {st.session_state.clip_model.model_name}")
    else:
        st.warning("请先加载模型")
    
    st.markdown("---")
    
    # 知识库管理
    st.header("3. 知识库管理")
    
    if st.session_state.clip_model is None:
        st.warning("请先加载CLIP模型")
        return
    
    # 创建标签页
    tab1, tab2, tab3 = st.tabs(["从文件夹构建", "上传图片添加", "管理现有知识库"])
    
    # 从文件夹构建知识库
    with tab1:
        st.subheader("从文件夹构建知识库")
        
        data_dir = "./data"
        products = get_products_from_data_dir(data_dir)
        
        if not products:
            st.warning("未找到可用的产品数据集。请检查数据目录。")
        else:
            selected_product = st.selectbox("选择产品", products)
            
            if selected_product:
                train_dir = os.path.join(data_dir, selected_product, "train")
                defect_types = get_defect_types_from_train_dir(train_dir)
                
                # 显示可用的缺陷类型
                st.markdown(f"**可用的缺陷类型**: {', '.join(defect_types)}")
                
                # 计算图片总数
                total_images = 0
                for defect_type in defect_types:
                    defect_dir = os.path.join(train_dir, defect_type)
                    for ext in ['png', 'jpg', 'jpeg']:
                        total_images += len(glob.glob(os.path.join(defect_dir, f"*.{ext}")))
                
                st.markdown(f"**总图片数**: {total_images}")
                
                # 构建知识库按钮
                if st.button("构建知识库", key="build_kb"):
                    with st.spinner(f"正在为 {selected_product} 构建知识库..."):
                        try:
                            # 构建知识库
                            start_time = time.time()
                            success = st.session_state.knowledge_base.build_from_folder(
                                product_name=selected_product,
                                data_dir=train_dir
                            )
                            end_time = time.time()
                            
                            if success:
                                st.success(f"{selected_product} 知识库构建完成！用时: {end_time - start_time:.2f}秒")
                            else:
                                st.error("知识库构建失败")
                        except Exception as e:
                            st.error(f"知识库构建错误: {str(e)}")
    
    # 上传图片添加到知识库
    with tab2:
        st.subheader("上传图片添加到知识库")
        
        # 选择或创建产品
        product_options = ["创建新产品"] + st.session_state.knowledge_base.get_product_list()
        product_choice = st.selectbox("选择产品", product_options, key="upload_product_choice")
        
        if product_choice == "创建新产品":
            new_product = st.text_input("输入新产品名称")
            if new_product:
                product_name = new_product
            else:
                st.warning("请输入产品名称")
                product_name = None
        else:
            product_name = product_choice
        
        # 设置缺陷类型
        defect_type = st.text_input("输入缺陷类型（例如: scratch, good 等）")
        
        # 文件上传器
        uploaded_files = st.file_uploader("上传图片", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
        
        if product_name and defect_type and uploaded_files:
            if st.button("添加到知识库", key="add_to_kb"):
                with st.spinner("处理上传的图片..."):
                    try:
                        # 创建临时目录保存上传的文件
                        upload_dir = os.path.join("./uploads", product_name, defect_type)
                        os.makedirs(upload_dir, exist_ok=True)
                        
                        # 保存上传的文件
                        saved_paths = []
                        for file in uploaded_files:
                            saved_path = save_upload_file(file, upload_dir)
                            saved_paths.append(saved_path)
                        
                        # 添加到知识库
                        if saved_paths:
                            success = st.session_state.knowledge_base.add_images(
                                product_name=product_name,
                                image_paths=saved_paths,
                                labels=[defect_type] * len(saved_paths)
                            )
                            
                            if success:
                                st.success(f"成功添加 {len(saved_paths)} 张图片到 {product_name} 知识库")
                            else:
                                st.error("添加图片失败")
                    except Exception as e:
                        st.error(f"处理上传图片错误: {str(e)}")
    
    # 管理现有知识库
    with tab3:
        st.subheader("管理现有知识库")
        
        # 获取知识库列表
        kb_list = st.session_state.knowledge_base.get_product_list()
        
        if not kb_list:
            st.info("暂无已创建的知识库")
        else:
            # 显示知识库列表
            kb_data = []
            for kb_name in kb_list:
                # 加载知识库元数据
                st.session_state.knowledge_base.load_or_create(kb_name)
                kb_info = st.session_state.knowledge_base.kb_info
                
                kb_data.append({
                    "知识库名称": kb_name,
                    "创建时间": kb_info.get("created_at", "未知"),
                    "图片数量": kb_info.get("total_images", 0),
                    "缺陷类型": ", ".join(kb_info.get("defect_types", []))
                })
            
            # 显示知识库表格
            kb_df = pd.DataFrame(kb_data)
            st.dataframe(kb_df, use_container_width=True)
            
            # 选择要管理的知识库
            selected_kb = st.selectbox("选择要管理的知识库", kb_list)
            
            if selected_kb:
                col1, col2 = st.columns(2)
                
                with col1:
                    if st.button("加载知识库", key="load_kb"):
                        with st.spinner(f"加载 {selected_kb} 知识库..."):
                            try:
                                success = st.session_state.knowledge_base.load_or_create(selected_kb)
                                if success:
                                    st.success(f"{selected_kb} 知识库加载成功")
                                    
                                    # 显示知识库信息
                                    st.markdown(f"**图片数量**: {st.session_state.knowledge_base.kb_info.get('total_images', 0)}")
                                    st.markdown(f"**缺陷类型**: {', '.join(st.session_state.knowledge_base.kb_info.get('defect_types', []))}")
                                else:
                                    st.error("知识库加载失败")
                            except Exception as e:
                                st.error(f"知识库加载错误: {str(e)}")
                
                with col2:
                    if st.button("删除知识库", key="delete_kb", type="primary"):
                        if st.session_state.knowledge_base:
                            confirm = st.checkbox("确认删除？此操作不可恢复", key="confirm_delete")
                            if confirm:
                                with st.spinner(f"删除 {selected_kb} 知识库..."):
                                    try:
                                        success = st.session_state.knowledge_base.delete_knowledge_base(selected_kb)
                                        if success:
                                            st.success(f"{selected_kb} 知识库已删除")
                                            st.experimental_rerun()
                                        else:
                                            st.error("知识库删除失败")
                                    except Exception as e:
                                        st.error(f"知识库删除错误: {str(e)}")
    
    st.markdown("---")
    
    st.markdown("""
    ### 知识库使用说明
    
    1. 先加载CN_CLIP模型
    2. 选择Milvus模式（Milvus Lite或服务器）
    3. 选择从文件夹构建知识库或上传图片添加到知识库
    4. 构建完成后可在"管理现有知识库"中查看和管理
    
    知识库构建过程会提取图像特征并保存为向量索引，以便后续快速检索相似图像。
    """) 